import os
from langchain_core.prompts import PromptTemplate, FewShotPromptTemplate
from langchain_experimental.tabular_synthetic_data.openai import create_openai_data_generator
from langchain_experimental.tabular_synthetic_data.prompts import SYNTHETIC_FEW_SHOT_PREFIX, SYNTHETIC_FEW_SHOT_SUFFIX
from langchain_openai import ChatOpenAI
from pydantic import SecretStr, BaseModel

# 读取API密钥
api_key = os.getenv("DASHSCOPE_API_KEY")
if not api_key:
    raise ValueError("请设置环境变量DASHSCOPE_API_KEY（阿里云百炼API-KEY）")

# 创建大语言模型实例
model = ChatOpenAI(
    model="qwen-plus-latest",
    temperature=0.5,
    max_tokens=None,
    timeout=None,
    max_retries=2,
    api_key=SecretStr(api_key),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)

# 1. 定义数据类型（Pydantic 2.x模型）
class MedicalBilling(BaseModel):
    patient_id: int
    patient_name: str
    diagnosis_code: str
    procedure_code: str
    total_charge: float
    insurance_claim_amount: float

# 2. 提供样例数据（简化格式）
examples = [
    {
        "MedicalBilling": "患者ID: 123456, 姓名: 张三, 诊断代码: J20.9, 程序代码: 99213, 总费用: 500.0, 保险索赔金额: 200.0"
    },
    {
        "MedicalBilling": "患者ID: 223456, 姓名: 李四, 诊断代码: J20.8, 程序代码: 88213, 总费用: 800.0, 保险索赔金额: 400.0"
    }
]

# 3. 创建提示模板（使用标准FewShotPromptTemplate）
openai_template = PromptTemplate(
    input_variables=["MedicalBilling"],
    template="{MedicalBilling}"
)

prompt_template = FewShotPromptTemplate(
    examples=examples,
    example_prompt=openai_template,
    prefix=SYNTHETIC_FEW_SHOT_PREFIX,
    suffix=SYNTHETIC_FEW_SHOT_SUFFIX,
    input_variables=["subject", "extra"]
)

# 调试：打印生成的提示内容
formatted_prompt = prompt_template.format(subject="医疗账单", extra="名字随机，且为中文，名字比较中二的")
print("生成的提示内容:")
print(formatted_prompt)

# 4. 创建数据生成器
generator = create_openai_data_generator(
    output_schema=MedicalBilling,
    llm=model,
    prompt=prompt_template,
)

# 5. 生成数据
result = generator.generate(
    subject="医疗账单",
    extra="名字随机，且为中文，名字比较中二的（例如：龙傲天、叶孤城）",
    runs=10,
)

print(result)